import * as dotenv from 'dotenv';
import { eq, inArray, isNull } from 'drizzle-orm';
import { drizzle } from 'drizzle-orm/node-postgres';
import { Pool } from 'pg';

import { chunks, fileChunks, messages } from '@/database/schemas';

// 确保环境变量已加载
dotenv.config();

console.log('开始清理孤立数据...');

try {
  // 创建数据库连接
  const pool = new Pool({
    connectionString: process.env.DATABASE_URL,
  });

  const db = drizzle(pool);

  // 1. 删除孤立消息（没有关联会话的消息）
  const orphanMessages = await db.delete(messages).where(isNull(messages.sessionId)).returning();

  console.log(`删除了 ${orphanMessages.length} 条孤立消息`);

  // 2. 删除孤立的文件块
  const orphanedChunks = await db
    .select({ chunkId: chunks.id })
    .from(chunks)
    .leftJoin(fileChunks, eq(chunks.id, fileChunks.chunkId))
    .where(isNull(fileChunks.fileId));

  const chunkIds = orphanedChunks.map((chunk) => chunk.chunkId);

  if (chunkIds.length > 0) {
    // 分批删除，每批500个
    const batchSize = 500;
    for (let i = 0; i < chunkIds.length; i += batchSize) {
      const batch = chunkIds.slice(i, i + batchSize);
      await db.delete(chunks).where(inArray(chunks.id, batch));
    }
  }

  console.log(`删除了 ${chunkIds.length} 个孤立文件块`);

  console.log('清理完成！');

  // 关闭数据库连接
  await pool.end();
} catch (error) {
  console.error('清理过程中发生错误:', error);
  throw error; // 抛出错误而不是使用 process.exit
}
